﻿#ifndef __ROAD_SIGN_MODEL_H__
#define __ROAD_SIGN_MODEL_H__
#include "base64.h"
#include "models/SoftMaxModel.h"

#include "JsColorModel.h"

using SoftMaxModel = darknet::SoftMaxModel;

class RoadSignModel : public SoftMaxModel
{
    using Options   = darknet::Options;
    using DataImage = darknet::DataImage;

public:
    static RoadSignModel *GetInstance()
    {
        static RoadSignModel s_instance;
        return &s_instance;
    }

public:
    RoadSignModel()
    {
        m_max_epochs      = 4000;
        m_batch           = 8;
        m_auto_stop_train = false;
        m_model_name      = "roadsigns";

#ifdef __EMSCRIPTEN__
        std::string root_directory = "/assets/" + GetModelName();
#else
        std::string root_directory = "D:/workspace/cpp/darknet/web/asset/assets/" + GetModelName();
#endif
        std::string model_directory   = root_directory + "/models";
        std::string model_filepath    = model_directory + "/" + GetDefaultModelFileName();
        std::string dataset_directory = root_directory + "/datasets";
        std::string dataset_filepath  = dataset_directory + "/dataset.txt";

        SetRootDirectory(root_directory);
        SetModelDirectory(model_directory);
        SetModelPath(model_filepath);
        SetDataSetDirectory(dataset_directory);
        SetDataSetPath(dataset_filepath);
#ifdef __EMSCRIPTEN__
        LoadDataSet();
        // std::cout << "load default dataset: " << GetDataSetPath() << " size = " << GetDataSet()->size() << std::endl;
#else
        std::filesystem::create_directories(root_directory);
        std::filesystem::create_directories(model_directory);
        std::filesystem::create_directories(dataset_directory);
#endif

        if (std::filesystem::exists(model_filepath) && LoadModelFromFile(model_filepath))
        {
            std::cout << "load prebuilt model from " << model_filepath << std::endl;
        }
        else
        {
            LoadNet();
        }

        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(左转)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(右转)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(直行)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(停车)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(隧道)"));
        CreateGetClassifyIndex(reinterpret_cast<const char*>(u8R"(禁止通行)"));
    }

    virtual void LoadNet()
    {
        if (GetNet()->IsLoaded()) return;

        auto net     = GetNet();
        auto options = std::make_shared<std::vector<std::shared_ptr<Options>>>();
        auto option  = std::make_shared<Options>("net");
        option->SetOption("width", 100);
        option->SetOption("height", 60);
        option->SetOption("channels", 3);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 32);
        options->push_back(option);

        option = std::make_shared<Options>("convolutional");
        option->SetOption("size", 3);
        option->SetOption("stride", 2);
        option->SetOption("filters", 16);
        options->push_back(option);

        option = std::make_shared<Options>("connected");
        option->SetOption("output", GetClassifySize());
        option->SetOption("activation", "linear");
        options->push_back(option);

        option = std::make_shared<Options>("softmax");
        options->push_back(option);

        net->Load(options);
        net->SetBatch(GetBatch());
    }

    bool LoadModelFromFile(const std::string &model_path)
    {
        return SoftMaxModel::LoadModelFromFile(model_path);
    }

    virtual void LoadDataSet()
    {
        SoftMaxModel::LoadDataSet();

#ifndef __EMSCRIPTEN__
        auto color_model = JsColorModel::GetInstance();
        auto dataset     = GetDataSet();

        for (auto it = dataset->begin(); it != dataset->end(); it++)
        {
            std::string image_path = it->first;

            auto img = std::make_shared<DataImage>();
            if (!img->Load(image_path))
            {
                std::cout << "invalid image path: " << image_path << std::endl;
                continue;
            }

            img = HandleInputImage(img);
            if (img == nullptr)
            {
                std::cout << "不存在红色路标!!!" << image_path << std::endl;
                continue;
            }

            auto path      = std::filesystem::path(image_path);
            auto filename  = path.filename();
            auto dirname   = path.parent_path().filename();
            auto directory = path.parent_path().parent_path();
            filename.replace_extension(".jpg");
            img->Save((directory / "auto" / (dirname.string() + "_" + filename.string())).string());
        }
#endif
    }

    virtual std::shared_ptr<DataImage> HandleInputImage(std::shared_ptr<DataImage> input_image)
    {
        m_x      = 0;
        m_y      = 0;
        m_width  = 0;
        m_height = 0;

        // 取中下部分做颜色检测, 避免背景影响
        auto input_image_width  = input_image->GetWidth();
        auto input_image_height = input_image->GetHeight();
        auto input_image_size   = 200;
        auto offset_x           = (input_image_width - input_image_size) / 2;
        auto offset_y           = input_image_height - input_image_size;
        input_image             = input_image->GetSubImage(offset_x, offset_y, input_image_size, input_image_size);
        // 提取颜色值
        auto color_model          = JsColorModel::GetInstance();
        auto roadsign_color_index = color_model->GetRedClassifyIndex();
        auto img_width            = input_image->GetWidth();
        auto img_height           = input_image->GetHeight();
        auto rects                = color_model->GetColorRect(input_image, roadsign_color_index, 0.2);
        if (rects == nullptr || rects[2] > 240 || rects[3] > 120) return nullptr;
        // std::cout << "width = " << rects[2] << ", height = " << rects[3] << std::endl;

        auto x = std::max(0, rects[0] - 10);
        auto y = std::max(0, rects[1] - 10);
        auto w = std::min(img_width, rects[2] + 20);
        auto h = std::min(img_height, rects[3] + 20);

        m_x      = offset_x + x;
        m_y      = offset_y + y;
        m_width  = w;
        m_height = h;

        input_image = input_image->GetSubImage(x, y, w, h);
        auto width  = input_image->GetWidth();
        auto height = input_image->GetHeight();

        if (width > GetNetWidth() || height > GetNetHeight()) input_image = input_image->ScaleFill(GetNetWidth(), GetNetHeight());

        for (int y = 0; y < height; y++)
        {
            for (int x = 0; x < width; x++)
            {
                auto r           = input_image->GetPixel(x, y, 0) * 0xff;
                auto g           = input_image->GetPixel(x, y, 1) * 0xff;
                auto b           = input_image->GetPixel(x, y, 2) * 0xff;
                auto color_index = color_model->Predict(r, g, b);
                auto color       = color_index == roadsign_color_index ? 1.0f : 0.0f;
                input_image->SetPixel(x, y, 0, color);
                input_image->SetPixel(x, y, 1, color);
                input_image->SetPixel(x, y, 2, color);
            }
        }
        return input_image;
    }

    inline int GetX() { return m_x; }
    inline int GetY() { return m_y; }
    inline int GetWidth() { return m_width; }
    inline int GetHeight() { return m_height; }

protected:
    int m_x;
    int m_y;
    int m_width;
    int m_height;
};

#endif

// 119,134,129